package com.example.demomcp.service;

import jakarta.annotation.Resource;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Component
public class DBToolService {

    @Resource
    private JdbcTemplate jdbcTemplate;

    private final String sql = "SELECT TABLE_NAME, TABLE_COMMENT FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'im_platform'";

    private final String schemaSql = "SELECT COLUMN_NAME, DATA_TYPE, COLUMN_COMMENT FROM INFORMATION_SCHEMA.COLUMNS " +
            "WHERE TABLE_SCHEMA = 'im_platform' AND TABLE_NAME = ?";

    @Tool(description = "获取所有可用的表名")
    public List<String> getTables() {
        List<Map<String, Object>> maps = jdbcTemplate.queryForList(sql);
        return maps.stream().map(map -> {
            String tableName = String.valueOf(map.get("TABLE_NAME"));
            String tableComment = String.valueOf(map.get("TABLE_COMMENT"));
            return tableName + " (" + tableComment + ")";
        }).collect(Collectors.toList());
    }

    @Tool(description = "根据表名获取Schema")
    public String getTableSchema(@ToolParam(description = "表名") List<String> tables) {
        return tables.stream().filter(t -> !t.isBlank()).map(tableName -> {
            List<Map<String, Object>> columns = jdbcTemplate.queryForList(schemaSql, tableName);
            String tablePrompt = columns.stream().map(map -> {
                String name = String.valueOf(map.get("COLUMN_NAME"));
                String type = String.valueOf(map.get("DATA_TYPE"));
                String comment = String.valueOf(map.get("COLUMN_COMMENT"));
                return String.format("%s (%s) - %s", name, type, comment);
            }).collect(Collectors.joining(", \n"));
            return String.format("Table: %s (%s)\n", tableName, tablePrompt);
        }).collect(Collectors.joining("\n"));
    }

    @Tool(description = "执行SQL查询结果")
    public List<Map<String, Object>> runSql(@ToolParam(description = "sql") String sql) {
        if (sql.contains("DELETE") || sql.contains("UPDATE") || sql.contains("INSERT")){
            throw new RuntimeException("执行SQL仅限于查询语句！");
        }
        return jdbcTemplate.queryForList(sql);
    }

}
